% This code simulates the model for EIS = 1. 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

function SIMDATA = Sim_EIS_1(l_q , l_q_burn , N_sim , nu , c_0 , node_n , Var_Names)

% INPUT:
% -- l_q: Number of sample periods
% -- l_q_burn: Number of burn-in periods
% -- N_sim: Number of simulations
% -- nu: constant gain parameter
% -- c_0: starting value of log real per capita consumption level
% -- Var_Names: names of variables to be stored
%
% OUTPUT:
% -- SIMDATA: stored simulations as a cell

%% 1. LOAD PARAMETERS
if nu == 0
    error('nu cannot be zero!')
end

I = param_EIS_1(nu);

Constraint = I.tmu_m + 0.5 * (sqrt(1 + I.nu) * (1 + I.nu * (I.lambda - 1)/I.alpha) - I.xi)^2 * I.sigma^2;

if Constraint >= 0
    error('Constraint violated!')
end
%% 2. SIMULATIONS
% additional parameters for simulation
T = l_q;                %Number of sample periods
T_burn = l_q_burn;            %Number of burn-in periods
T_total = T_burn + T;     %Total periods including burn-in periods for simulation
d_0 = c_0 + I.d_c_mean;   %Starting value of log dividend

%% 2.1 Simulate the economy
% Simulate quarterly iid consumption growth, column for each simulation
cg_full = normrnd(I.mu , I.sigma , T_total , N_sim); 

% Simulate quarterly shocks to dividend growth
dg_shock = normrnd(0 , I.sigma_d , T_total , N_sim); 

% Log real consumption level from 1 to T_total
c_full = c_0 + cumsum(cg_full); 

% Log real dividend level from 1 to T_total
d_full = zeros(T_total , N_sim); 

d_full(1 , :) = I.lambda * cg_full(1 , :) + (1 - I.alpha) * d_0 + I.alpha * c_0 + I.alpha * I.mu_dc + dg_shock(1 , :);

for k = 2 : T_total
  d_full(k , :) = I.lambda * cg_full(k , :) + (1 - I.alpha) * d_full(k - 1 , :) + I.alpha * c_full(k - 1 , :) + I.alpha * I.mu_dc + dg_shock(k , :);
end

% Log dividend growth from 1 to T_total
dg_full = [d_full(1 , :) - d_0 ; d_full(2 : end , :) - d_full(1 : end - 1 , :)]; 

% Four-quarter change
cg_yoy_full = (c_full(5 : end , :) - c_full(1 : end - 4 , :))/4;
dg_yoy_full = (d_full(5 : end , :) - d_full(1 : end - 4 , :))/4;

% Construct State Variable \tilde{\mu}
tmu_full = expweightedavg(cg_full , I.nu);
if I.nu == 0
    tmu_full = I.mu * ones(T_total , N_sim);
end

% Calculate Risk-free rate from 1 to T_total
rf_full = -I.tmu_m + tmu_full - 0.5 * I.xi^2 * I.sigma^2; 

Rf_full = exp(rf_full) - 1; 

%% 2.2 Calculate price
P_full = zeros(T_total , N_sim);                 %Price level

% Max number of terms to calculate explicitly
J = round(log(0.0001)/log(1 - I.alpha)); 

% Calculate the first J terms of A_{c,n}: A_{c,1},...,A_{c,J} as in Equation (D.17)
A_n_inc = ( sqrt(1+I.nu) ... 
          * ( I.nu * (I.lambda-1) * (1 - (1 - I.alpha).^(0 : J - 1))/I.alpha ...
          + (I.lambda - 1) * (1-I.alpha).^(0 : J - 1) + 1 ) ...
          - I.xi ).^2;
      
A_N = cumsum(A_n_inc); % These are the A_{c,n}

% Further incremental terms to A_{c,n} will be approximated by A_n_inc_lim
A_n_inc_lim = ( sqrt(1 + I.nu) * ( I.nu * (I.lambda - 1)/I.alpha + 1) - I.xi )^2;

% Calculate the first J dividend strips' prices at time t. first calculate
% the multiplicative term that is invariant to time on the RHS of Equation (D.16)
p_cons = ( I.tmu_m * (1 : J) ... % n\tmu_m
       + (1 - (1 - I.alpha).^(1 : J)) * I.mu_dc ...%(1 - (1-\alpha)^n)\mu_{dc}
       + 0.5 * A_N * I.sigma^2 ... %0.5 * A_{c,n} \sigma^2
       + 0.5 * (1 - (1 - I.alpha).^(2 * (1 : J))) * I.sigma_d^2/(1 - (1 - I.alpha)^2) )'; % 0.5 * A_{d,n} \sigma_d^2

% Multiplier that is invariant with respect to simulations or time in
% Equation (D.31) and (D.32)
approx_cons = exp( I.mu_dc + 0.5 * I.sigma_d^2/(1 - (1 - I.alpha)^2) ...
            + (J + 1) * I.tmu_m + 0.5 * A_N(J) * I.sigma^2 + 0.5 * A_n_inc_lim * I.sigma^2 ) ...%V_J
            /(1 - exp(I.tmu_m + 0.5 * A_n_inc_lim * I.sigma^2));

% Calculate total price using Equation (D.31)        
for n = 1 : T_total 
    %Calculate p_t^n for maturity up to J
    p_t_n_temp = d_full(n , :) ...  
  + bsxfun(@times , ( 1 - (1 - I.alpha).^(1 : J) )' , (c_full(n , :) - d_full(n , :) + tmu_full(n,:) * (I.lambda - 1)/I.alpha) ) ...
  + p_cons;

    %Calculate aggregate price
    P_full(n,:) = sum(exp(p_t_n_temp) , 1) ... % First J dividend strip
                + exp(c_full(n , :)) .* exp( (I.lambda - 1) * tmu_full(n,:)/I.alpha ) * approx_cons;    
end

%Calculate trailing 4-quarter sum of dividends
D_ttm_full = movsum(exp(d_full) , 4 , 'Endpoints' , 'discard');

%--------------------------------------------------------------------------
% %The following code calculates the analytical price when alpha = 1.
% Compare to make sure algorithm is right

% Const_1 = I.tmu_m + 0.5 * (sqrt(1 + I.nu) * I.lambda - I.xi)^2 * I.sigma^2;
% Const_2 = I.tmu_m + 0.5 * (sqrt(1 + I.nu) * (I.nu * (I.lambda - 1) + 1) - I.xi)^2 * I.sigma^2;
% 
% P_alpha1 = exp(c_full + (I.lambda - 1) * tmu_full + I.mu_dc + 0.5 * I.sigma_d^2) * exp(Const_1)/(1-exp(Const_2));
%-------------------------------------------------------------------------- 

%% 2.3 Calculate other quantities
% P/D ratio
PD_ttm_full = P_full(4 : end , :) ./ D_ttm_full; 
pd_ttm_full = log(PD_ttm_full);

PD_full = P_full./exp(d_full);
pd_full = log(PD_full);

% Objective realized returns
ret_full = log( P_full(2 : end , :) + exp(d_full(2 : end , :)) ) ... % log(P_{t+1} + D_{t+1})
         - log( P_full(1 : end - 1 , :) ); % - log(P_t)
     
ret_MS_full = movsum(ret_full , 4 , 'Endpoints' , 'discard');     
     
Ret_full = exp(ret_full) - 1;

EXRET_full = exp(ret_full) - exp(rf_full(1 : end - 1 , :)); % R_{t+1} - R_{f,t}

exret_full = ret_full - rf_full(1:end-1,:); %r_{t+1} - r_{f,t}

% Experienced dividend growth
tmud_full = expweightedavg(dg_full , I.nu);

% Experience returns
tmur_1_full = expweightedavg(ret_full , I.nu);   %2 to T_total, use experienced log returns(in the paper)
tmur_2_full = expweightedavg(exret_full , I.nu); %2 to T_total, use experienced log excess returns(in case)

%% 3. SIMULATED DATA

%% 3.1 Drop burn-in periods
% All data are re-labled from 1 to T as realizations. 
% Be careful with risk-free rate

% Endowment Process
cg = cg_full(end - T + 1 : end , :);
dg = dg_full(end - T + 1 : end , :);
c  =  c_full(end - T + 1 : end , :);
d  =  d_full(end - T + 1 : end , :);
cg_yoy = cg_yoy_full(end - T + 1 : end , :);
dg_yoy = dg_yoy_full(end - T + 1 : end , :);

% State variable proxies
tmu = tmu_full(end - T + 1 : end , :);
tmud = tmud_full(end - T + 1 : end , :);
tmur_1 = tmur_1_full(end - T + 1 : end , :);
tmur_2 = tmur_2_full(end - T + 1 : end , :);
tmur = tmur_1; % For now we use the experienced log returns specification

% Price Quantity
rf = rf_full(end - T + 1 : end , :);
Rf = Rf_full(end - T + 1 : end , :);

P = P_full(end - T + 1 : end , :);
D_ttm = D_ttm_full(end - T + 1 : end , :);
pd_ttm = pd_ttm_full(end - T + 1 : end , :);
pd = pd_full(end - T + 1 : end , :);

ret = ret_full(end - T + 1 : end , :);
ret_MS = ret_MS_full(end - T + 1 : end , :);
Ret = Ret_full(end - T + 1 : end , :);
EXRET = EXRET_full(end - T + 1 : end , :);
exret = exret_full(end - T + 1 : end , :);

%% 3.2 Subjective growth expectation

% Here we calculate expectations of future average 1-quarter and 20-quarter
% growth rates
t_fc = 20; % Number of periods to forecast
LTG = nan(T , N_sim); % Average subjective dividend growth expectations for the next 5 years
Y1G = nan(T , N_sim); % Average subjective dividend growth expectations for the next 1 year
for k = 1 : T
    % Store the next 20-quarter subjective dividend growth expectations at each point of time
    sub_dg_temp = repmat(tmu(k , :) , t_fc , 1) .* repmat(1 + (I.lambda - 1) * (1 - I.alpha).^(0 : t_fc - 1)' , 1 , N_sim) ...
                - repmat(d(k , :) - c(k , :) - I.mu_dc , t_fc , 1) .* repmat(I.alpha * (1 - I.alpha).^(0 : t_fc - 1)' , 1 , N_sim);
            
    LTG(k , :) = mean(sub_dg_temp); 
    Y1G(k , :) = mean(sub_dg_temp(1 : 4 , :));
end

%% 3.3 Solve for subjective equity premium using the projection method
% Follow the method in Pohl, Schmedders, Wilms(2018 JF)

%--------------------------------------------------------------------------
% Generate Grid and Basis Functions
%--------------------------------------------------------------------------
% Two-Dimensional State Space. 4 s.d. as in Pohl et al.(2018)
mu_min = I.mu - 4 * mean(std(tmu),2);
mu_max = I.mu + 4 * mean(std(tmu),2);

mu_dc_min = I.d_c_mean - 4 * I.d_c_mean_sig;
mu_dc_max = I.d_c_mean + 4 * I.d_c_mean_sig;

% Another state variable
d_c_diff = d - c;

v = 200;     %Number of points for calculating integral

% Use SQP algorithm 
options_fmincon = optimoptions('fmincon','Algorithm','sqp','MaxFunEvals',1000,'TolFun',1e-3,'Display','off');

% Generate polynomials
fspace_tmudc = fundefn('cheb',node_n,[mu_min mu_dc_min], [mu_max mu_dc_max]);          
tmu_dc_grid = funnode(fspace_tmudc); 
beta_PD_0 = zeros(node_n(1) * node_n(2),1);                       
[beta_PD_cointegration , ~ , exitflag] = ...
fmincon(@(gamma)fun_fmincon_const(gamma),beta_PD_0,[],[],[],[],[],[],@(beta)projection_constraint_PD_EIS_1(beta, I , v, fspace_tmudc, tmu_dc_grid),options_fmincon);

pd_proj = zeros(T,N_sim);

for i = 1 : N_sim
  grid_PD_sim_temp{1 , 1} = tmu(: , i);
  grid_PD_sim_temp{1 , 2} = d_c_diff(: , i);
  fbasis_sim_temp = funbas(fspace_tmudc , grid_PD_sim_temp); 
  pd_proj_temp = fbasis_sim_temp * beta_PD_cointegration;
  pd_proj(: , i) = pd_proj_temp(1 : T+1 : T^2);
end

% RMSE of analytical P/D and projection P/D
RMSE_PD = sqrt(sum((pd - pd_proj).^2)/T);

%--------------------------------------------------------------------------
% Calculate the subjective equity premium
%--------------------------------------------------------------------------
sub_Ret = zeros(T , N_sim); 
for i = 1 : N_sim
    for k = 1 : T
       sub_Ret(k,i) = calculate_subRet_proj(d_c_diff(k,i), tmu(k,i), beta_PD_cointegration, I , v, fspace_tmudc) - 1; 
    end
end

% Note the subjective expected returns are known at time t, thus subtract
% the contemporaneous risk-free rate
sub_ExRet = 1 + sub_Ret - exp(rf);

sub_ret = log(1 + sub_Ret); %log subjectively expected returns
sub_exret = sub_ret - rf; 

%% 4. STORE SIMULATIONS
SIMDATA = cell(length(Var_Names) , 1);
for i = 1 : length(Var_Names)
    SIMDATA{i} = eval(Var_Names{i});
end

end


               